"""
🔥 LangGraph 子图 (Subgraph) 实战示例：日志分析主图 + 两个子图

-----------------------------------------
一、核心概念回顾
-----------------------------------------
📘 子图（Subgraph）
LangGraph 中的图（Graph）由节点（Node）+ 边（Edge）组成，
每个节点就是一个可执行函数，边表示执行顺序。

当流程变复杂时，我们可以把某一段逻辑单独拆成「子图」。
子图本质上就是一个能独立运行的小 StateGraph，
也可以嵌入主图作为一个节点执行。

优点：
✅ 模块化（逻辑独立、好维护）
✅ 可复用（同一子图可在不同主图使用）
✅ 可调试（可以单独运行）
✅ 可组合（子图中还能嵌子图）

本例主图包含两个子图：
- failure_analysis：失败日志分析
- question_summarization：问题总结与报告生成
"""

from operator import add
from typing import TypedDict, Optional, List, Annotated, Dict

from langgraph.constants import START, END
from langgraph.graph import StateGraph


# =========================================================
# 一、定义基础结构：日志（Logs）
# =========================================================
class Logs(TypedDict):
    """单条日志结构"""
    id: str
    question: str
    answer: str
    docs: Optional[List]      # 检索到的文档
    grade: Optional[int]      # 评分（有分表示失败）
    grader: Optional[str]     # 评分者
    feedback: Optional[str]   # 反馈文本


# =========================================================
# 二、定义子图 1：失败分析 (Failure Analysis)
# =========================================================
class FailureAnalysisState(TypedDict):
    """失败分析子图状态结构"""
    docs: List[Logs]
    failures: List[Logs]
    fa_summary: str


def get_failures(state):
    """筛选出失败日志（带有 grade 的认为失败）"""
    docs = state["docs"]
    failures = [doc for doc in docs if "grade" in doc]
    print(f"📊 [FailureAnalysis] 发现失败日志数量: {len(failures)}")
    return {"failures": failures}


def generate_fa_summary(state):
    """生成失败分析总结（此处模拟，可接入 LLM）"""
    failures = state["failures"]
    fa_summary = f"共发现 {len(failures)} 条失败日志，主要问题为：文档检索质量差"
    print(f"🧠 [FailureAnalysis] 总结完成: {fa_summary}")
    return {"fa_summary": fa_summary}


# 构建子图
fa_builder = StateGraph(FailureAnalysisState)
fa_builder.add_node("get_failures", get_failures)
fa_builder.add_node("generate_fa_summary", generate_fa_summary)

fa_builder.add_edge(START, "get_failures")
fa_builder.add_edge("get_failures", "generate_fa_summary")
fa_builder.add_edge("generate_fa_summary", END)


# =========================================================
# 三、定义子图 2：问题总结 (Question Summarization)
# =========================================================
class QuestionSummarizationState(TypedDict):
    """问题总结子图状态结构"""
    docs: List[Logs]
    qs_summary: str
    report: str


def generate_qs_summary(state):
    """生成问题总结"""
    docs = state["docs"]
    qs_summary = f"共收到 {len(docs)} 条日志，涉及多个问题主题"
    print(f"📝 [QuestionSummary] 问题总结: {qs_summary}")
    return {"qs_summary": qs_summary}


def generate_report(state):
    """生成报告（此处模拟）"""
    qs_summary = state["qs_summary"]
    report = f"[报告] 自动生成：{qs_summary}。整体表现一般。"
    print(f"📢 [QuestionSummary] 生成报告: {report}")
    return {"report": report}


# 构建子图
qs_builder = StateGraph(QuestionSummarizationState)
qs_builder.add_node("generate_qs_summary", generate_qs_summary)
qs_builder.add_node("generate_report", generate_report)

qs_builder.add_edge(START, "generate_qs_summary")
qs_builder.add_edge("generate_qs_summary", "generate_report")
qs_builder.add_edge("generate_report", END)


# =========================================================
# 四、主图 (EntryGraph)
# =========================================================
class EntryGraphState(TypedDict):
    raw_logs: Annotated[List[Dict], add]
    docs: Annotated[List[Logs], add]
    fa_summary: str
    report: str


# 模拟数据
question_answer = Logs(
    id="1",
    question="如何导入 ChatOpenAI?",
    answer="使用 from langchain_openai import ChatOpenAI"
)

question_answer_feedback = Logs(
    id="2",
    question="如何使用 Chroma 向量存储?",
    answer="可以用 from langchain_community.vectorstores import Chroma",
    grade=2,
    grader="系统",
    feedback="回答未聚焦 Chroma 的使用细节"
)


def convert_logs_to_docs(state):
    """将原始日志转换为标准 Logs 结构"""
    print("🔄 [Main] 转换原始日志为标准结构")
    docs = [question_answer, question_answer_feedback]
    return {"docs": docs}


# 创建主图
entry_builder = StateGraph(EntryGraphState)
entry_builder.add_node("convert_logs_to_docs", convert_logs_to_docs)

# 嵌入两个子图
entry_builder.add_node("failure_analysis", fa_builder.compile())
entry_builder.add_node("question_summarization", qs_builder.compile())

# 定义执行流程
entry_builder.add_edge(START, "convert_logs_to_docs")
entry_builder.add_edge("convert_logs_to_docs", "failure_analysis")
entry_builder.add_edge("convert_logs_to_docs", "question_summarization")
entry_builder.add_edge("failure_analysis", END)
entry_builder.add_edge("question_summarization", END)

# 编译主图
app = entry_builder.compile()


# =========================================================
# 五、运行与结果可视化
# =========================================================
if __name__ == "__main__":
    # 执行主图
    print("\n🚀 开始运行主图 EntryGraph...\n")
    final_state = app.invoke({"raw_logs": []})

    print("\n✅ 运行完成！最终状态输出如下：")
    for key, value in final_state.items():
        print(f"  {key}: {value}")

    # 生成可视化图
    graph_png = app.get_graph(xray=1).draw_mermaid_png()
    with open("langgraph_subgraph_demo.png", "wb") as f:
        f.write(graph_png)
    print("\n📈 流程图已生成: langgraph_subgraph_demo.png")
